import numpy as np
import sim_utils as ut

#### 0.read data
imgs = np.load('small_dataset_imgs_1000.npy')
labels = np.load('small_dataset_labels_1000.npy')

#### 1. load weight
f1_filters      = 32
f1_kernel_size  = 3
f1_stride       = 1
f1_padding      = 0
f1_feature_size = 26

f2_filters      = 32
f2_kernel_size  = 3
f2_stride       = 2
f2_padding      = 0
f2_feature_size = 6

fc1_out_ch = 54
fc2_out_ch = 10

f1_q_mapping = np.load('conv1.weight_quantized.npy').reshape(f1_filters, f1_kernel_size ** 2).transpose(1, 0)
f2_q_mapping = np.load('conv2.weight_quantized.npy').reshape(f2_filters, f2_kernel_size ** 2 * f1_filters).transpose(1, 0)
fc1_q_mapping = np.load('fc1.weight_quantized.npy').transpose([1,0])
fc2_q_mapping = np.load('fc2.weight_quantized.npy').transpose([1,0])
##################################1. write weight ################################
                		        ######         ######
                		        ######         ######
                		        ######         ######
                		        ######         ######
                		     ############   ############
                		      ##########     ##########
                		       ########       ########
                		        ######         ######
                		         ###            ###
                		          #              #
###################################################################################
'''
dic =  {0:"0",1:"1",2:"2",3:"3",4:"4",5:"5",6:"6",7:"7",-1:"f" ,-2:"e" ,-3:"d" ,-4:"c" ,-5:"b" ,-6:"a" ,-7:"9",-8:"8"}
dic2 = {0:"0",1:"1",2:"2",3:"3",4:"4",5:"5",6:"6",7:"7",-1:"15",-2:"14",-3:"13",-4:"12",-5:"11",-6:"10",-7:"9",-8:"8"}
mapping_weight = np.zeros([576,128],np.int8)
for i in range(0,9):
    for j in range(0,32):
        mapping_weight[i,j] = f1_q_mapping[i,j]
for i in range(0,288):
    for j in range(32,64):
        mapping_weight[i,j] = f2_q_mapping[i,j-32]
for i in range(0,288):
    for j in range(64,64+54):
        mapping_weight[i,j] = fc1_q_mapping[i,j-64-54]
for i in range(0,54):
    for j in range(118,128):
        mapping_weight[i,j] = f2_q_mapping[i,j-118]
# weight file_data for sim
file_data = ""
for i in range(576):
    for j in range(128):
        file_data+= dic2[mapping_weight[i,j]]+","
    file_data+="\n"
with open("cnn_weight_rom.txt","w",encoding="utf-8") as f:
    f.write(file_data)
# weight file_data for fpga
file_data = "memory_initialization_radix=16;\nmemory_initialization_vector=\n"
for i in range(576):
    for j in range(128):
        file_data+= dic[mapping_weight[i,j]]
    file_data+=",\n"
with open("cnn_weight_rom.coe","w",encoding="utf-8") as f:
    f.write(file_data)
'''

##################################2. begin inference ##############################
                		        ######         ######
                		        ######         ######
                		        ######         ######
                		        ######         ######
                		     ############   ############
                		      ##########     ##########
                		       ########       ########
                		        ######         ######
                		         ###            ###
                		          #              #
##################################################################################
index       = 0
total_img   = 1
correct_img = 0
for index in range(total_img):
    img     = imgs[index]
    label   = labels[index]
    img_quan = np.round(img * 7)
    # first, img2col->mm->col2img->relu->pooling
    f1_input = ut.feature_map_to_input(img_quan,stride=f1_stride,kernel_size=f1_kernel_size,padding=f1_padding)
    f1_out_col = np.dot( f1_input.T ,f1_q_mapping )
#    print("f1_out_col:", f1_out_col.min(), f1_out_col.max())
    f1_out_img = ut.output_to_feature_map(f1_out_col,f1_feature_size,f1_feature_size)
    f1_out_img[f1_out_img<0] = 0
    f1_pooling = ut.pooling(f1_out_img,2)
    # second layer , img2col->mm->col2img->relu->pooling
    f2_input = ut.feature_map_to_input(f1_pooling,kernel_size=f2_kernel_size,stride=f2_stride,padding=f2_padding)
    f2_out_col = np.dot(f2_input.T,f2_q_mapping)//16
#    print("f2_out_col:",f2_out_col.min(),f2_out_col.max())
    f2_out_img = ut.output_to_feature_map(f2_out_col,f2_feature_size,f2_feature_size)
    f2_out_img[f2_out_img<0] = 0
    f2_pooling = ut.pooling(f2_out_img,2)
    # fc1 , mm->relu
    fc1_out = np.dot(f2_pooling.reshape(1,-1),fc1_q_mapping)//16
#    print("fc1_out:", fc1_out.min(), fc1_out.max())
    fc1_out[fc1_out<0] = 0
    # fc2  mm-> max
    fc2_out = np.dot(fc1_out.reshape(1,-1),fc2_q_mapping)//16
#    print("fc2_out:", fc2_out.min(), fc2_out.max())
    predict = fc2_out.argmax()
    if(predict==label):
        correct_img += 1
print(fc2_out)
print(correct_img/total_img)
